# -*- coding: utf-8 -*-
# @Author: Cursor
# @Date: 2025-01-04
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-01-04

# @Description: document processor for simple pdf ingestion

from datetime import datetime
from typing import Optional  # Add this import

import logging
from core.exception import CustomException

from config.settings import *

from crewplus.apps.rag.schemas.ingest_status import IngestStatus

from crewplus.apps.rag.processors.base_processor import DocumentProcessor
from crewplus.apps.rag.schemas import Document  
from crewplus.apps.rag.schemas.ingest_request import IngestRequest

from crewplus.services.vdb_service import VDBService
from crewplus.apps.rag.utils.document_util import unify_documents_meta 
from crewplus.apps.rag.loaders.pdf_docintel_loader import PDFDocintelLoader
from crewplus.apps.rag.loaders.excel_document_loader import ExcelDocumentLoader
from langchain_core.vectorstores.base import VectorStore
from langchain_core.embeddings import Embeddings
from langchain_milvus import Zilliz

from langchain_community.document_loaders import UnstructuredExcelLoader

class ExcelDocumentProcessor(DocumentProcessor):
    async def process(self, message: IngestRequest, collection_name: str, vector_store: Optional[VectorStore] = None, embeddings: Optional[Embeddings] = None) -> Document:
        # Implement the logic to process PDF documents
        document = Document()  # Create a new Document instance
        logging.info("ingesting excel -- start -- " + message.url)
        # Relate this vectorized Document to the knowledge base
        document.source_url = message.url
        document.kbase_id = message.kbase_id
        document.source_type = message.source_type
        document.file_type = message.file_type
        document.title = message.title
        
        # get start UTC datetime
        document.ingestion_start_time = str(datetime.now())
        
        vdbsrv = VDBService()
        embeddings = vdbsrv.get_embeddings()
        milvus_store: Zilliz = vdbsrv.get_vector_store(collection_name, embeddings)
                        
        vdbsrv.delete_old_indexes(message.url, milvus_store)
        
        # TODO: fill summary and content
        
        try:
            # load pdf
            file_path = message.url

            # use different parser for different type of pdf
            # if message.parser == 'docintel':
            #     loader = PDFDocintelLoader(file_path)
            #     pages = loader.load()
            # else:
            #     loader = PyPDFLoader(file_path)
            #     pages = loader.load_and_split()
            loader = UnstructuredExcelLoader(file_path)
            data = loader.load()
            # logging.info("加载excel by ExcelDocumentLoader")
            # data=ExcelDocumentLoader(file_path).lazy_load()
            # logging.info("加载excel by ExcelDocumentLoader end")
            udocs = unify_documents_meta(data, file_type=message.file_type, title=message.title)

            milvus_store.from_documents(
                udocs,
                embedding=embeddings,
                collection_name=collection_name,
                connection_args=MILVUS_CONNECTION_ARGS
            )          
            
            # get end UTC datetime
            document.ingestion_end_time = str(datetime.now())       
            
            document.ingestion_status = IngestStatus.INGESTED        
        except Exception as e:
            document.ingestion_status = IngestStatus.NOT_INGESTED
            
            raise CustomException(str(e), code=IngestStatus.NOT_INGESTED)
        
        return document